import os

import yaml
import pandas as pd
from volcenginesdkarkruntime import Ark

# 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
# 初始化Ark客户端，从环境变量中读取您的API Key
client = Ark(
    # 此为默认路径，您可根据业务所在地域进行配置
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    api_key=os.environ.get("ARK_API_KEY"),
)

def ask_llm(role, content):
    print("----- standard request -----")
    ask_str = f"{content}属于居民消费八大类中的哪一个分类？"
    completion = client.chat.completions.create(
    # 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
        model="doubao-1-5-thinking-pro-250415",
        messages=[
            {"role": "system", "content": role},
            {"role": "user", "content": ask_str},
        ],
    )
    print(completion.choices[0].message.content)
    return completion.choices[0].message.content


df = pd.read_excel('2020产品部门对应消费八大类.xlsx', index_col=0)

dic = {}
for i in range(len(df)):
    print(i)
    role = "你是一个精通投入产出和居民消费八大类数据的助手，你需要根据用户的问题，回答用户的问题。"
    content = df.index[i]
    req = ask_llm(role, content)
    dic[content] = req
    
df_req = pd.Series(dic)
df_req.to_excel("豆包问每个行业的分类.xlsx")

# # 写入 YAML 文件
# with open('data.yaml', 'w') as f:
#     yaml.dump(dic, f, sort_keys=False)

# # 读取 YAML 文件
# with open('data.yaml', 'r') as f:
#     loaded_data = yaml.safe_load(f)